import torch

torch.set_default_dtype(torch.float64)
torch.set_default_tensor_type(torch.DoubleTensor)

RESTART_NUMS = 20
ROUND_NUMS = 1
EPOCHS_NUMS = 500

case_name = "zhrR8"
DIM = 17  # 系统输入维度
N_H = 1  # 隐藏层层数
D_H = 20  # 每个隐藏层包含的神经元个数
VERBOSE = True  # 是否打印输出
FINE_TUNE = True  # false: 初始化参数, true: 微调参数，无需初始化参数
USE_CDINN = True  # false: nn, true: clamp(weight)
# USE_CDINN = False  # false: nn, true: clamp(weight)

alpha = 0.1  # B(f(x))-(1-alpha)B(x)
time_step = 0.1  # x_{t+1} = x_t + time_step * x'
piece_len = 0.3

BATCH_SIZE = 16  # mini_batch_train_size
# piece_num = [[1, 2, 2, 2, 1, 1, 2, 2, 2, 1, 1, 2, 2, 2, 2], [1, 2, 2, 2, 1, 1, 2, 2, 2, 1, 1, 2, 2, 2, 2], [1, 2, 2, 2, 1, 1, 2, 2, 2, 1, 1, 2, 2, 2, 2]]  # generate data, init/unsafe/domain
piece_num = [1024, 1024, 1024]  # generate data, init/unsafe/domain
cex_num = [[1, 1, 1, 2, 1, 2, 1, 1, 2, 1, 2, 1, 1, 2, 1, 2, 1], [1, 1, 1, 2, 1, 2, 1, 1, 2, 1, 2, 1, 1, 2, 1, 2, 1], [1, 1, 1, 2, 1, 2, 1, 1, 2, 1, 2, 1, 1, 2, 1, 2, 1]]  # add cex, init/unsafe/domain
cex_step = 0.001  # add_cex_step

TOL_INIT = 0.0  # loss = relu(B(x) + tol)
TOL_UNSAFE = 0.0
TOL_DOMAIN = 0.0

DECAY_DOMAIN = 1  # p1 * init + p2 * unsafe + p3 * domain
DECAY_INIT = 1
DECAY_UNSAFE = 1

LOSS_OPT_FLAG = 1e-16  # epoch_loss < flag ==> train_success
TOL_DATA_GEN = 1e-16  # -1.8 + tol <= x0 <= 2 + tol; -0.5 + tol <= x0 <= 0.5 + tol

init_min = [-0.3] + [-0.2] * 16
init_max = [0.0] + [0.3] * 16
init_shape = 1

unsafe_min = [-0.2] + [-0.3] * 16
unsafe_max = [-0.15] + [-0.25] * 16
unsafe_shape = 1

domain_min = [-0.3] * 17
domain_max = [0.3] * 17


def is_unsafe(xt):
    return unsafe_min[0] <= xt[0] <= unsafe_max[0] and unsafe_min[1] <= xt[1] <= unsafe_max[1] \
        and unsafe_min[2] <= xt[2] <= unsafe_max[2] and unsafe_min[3] <= xt[3] <= unsafe_max[3] \
        and unsafe_min[4] <= xt[4] <= unsafe_max[4] and unsafe_min[5] <= xt[5] <= unsafe_max[5] \
        and unsafe_min[6] <= xt[6] <= unsafe_max[6] and unsafe_min[7] <= xt[7] <= unsafe_max[7] \
        and unsafe_min[8] <= xt[8] <= unsafe_max[8] and unsafe_min[9] <= xt[9] <= unsafe_max[9] \
        and unsafe_min[10] <= xt[10] <= unsafe_max[10] and unsafe_min[11] <= xt[11] <= unsafe_max[11] \
        and unsafe_min[12] <= xt[12] <= unsafe_max[12] and unsafe_min[13] <= xt[13] <= unsafe_max[13] \
        and unsafe_min[14] <= xt[14] <= unsafe_max[14] and unsafe_min[15] <= xt[15] <= unsafe_max[15] \
        and unsafe_min[16] <= xt[16] <= unsafe_max[16]


def is_init(xt):
    return init_min[0] <= xt[0] <= init_max[0] and init_min[1] <= xt[1] <= init_max[1] \
        and init_min[2] <= xt[2] <= init_max[2] and init_min[3] <= xt[3] <= init_max[3] \
        and init_min[4] <= xt[4] <= init_max[4] and init_min[5] <= xt[5] <= init_max[5] \
        and init_min[6] <= xt[6] <= init_max[6] and init_min[7] <= xt[7] <= init_max[7] \
        and init_min[8] <= xt[8] <= init_max[8] and init_min[9] <= xt[9] <= init_max[9] \
        and init_min[10] <= xt[10] <= init_max[10] and init_min[11] <= xt[11] <= init_max[11] \
        and init_min[12] <= xt[12] <= init_max[12] and init_min[13] <= xt[13] <= init_max[13] \
        and init_min[14] <= xt[14] <= init_max[14] and init_min[15] <= xt[15] <= init_max[15] \
        and init_min[16] <= xt[16] <= init_max[16]


def is_domain(xt):
    return domain_min[0] <= xt[0] <= domain_max[0] and domain_min[1] <= xt[1] <= domain_max[1] \
        and domain_min[2] <= xt[2] <= domain_max[2] and domain_min[3] <= xt[3] <= domain_max[3] \
        and domain_min[4] <= xt[4] <= domain_max[4] and domain_min[5] <= xt[5] <= domain_max[5] \
        and domain_min[6] <= xt[6] <= domain_max[6] and domain_min[7] <= xt[7] <= domain_max[7] \
        and domain_min[8] <= xt[8] <= domain_max[8] and domain_min[9] <= xt[9] <= domain_max[9] \
        and domain_min[10] <= xt[10] <= domain_max[10] and domain_min[11] <= xt[11] <= domain_max[11] \
        and domain_min[12] <= xt[12] <= domain_max[12] and domain_min[13] <= xt[13] <= domain_max[13] \
        and domain_min[14] <= xt[14] <= domain_max[14] and domain_min[15] <= xt[15] <= domain_max[15] \
        and domain_min[16] <= xt[16] <= domain_max[16]


def vector_field(x):
    def f(i, x):
        if i == 1:
            return time_step * ((x[:, 1] + 2 * x[:, 2] + 2 * x[:, 3] + 2 * x[:, 4] + 2 * x[:, 5] + 2 * x[:, 6] + 2 * x[:, 7] + 2 * x[:, 8] + x[:, 9]) / 100.0 + 1) + x[:, 0]
        elif i == 2:
            return time_step * (x[:, 2]) + x[:, 1]
        elif i == 3:
            return time_step * (-10 * (x[:, 1] - x[:, 1] * x[:, 1] * x[:, 1] / 6.0) - x[:, 1]) + x[:, 2]
        elif i == 4:
            return time_step * (x[:, 4]) + x[:, 3]
        elif i == 5:
            return time_step * (-10 * (x[:, 3] - x[:, 3] * x[:, 3] * x[:, 3] / 6.0) - x[:, 1]) + x[:, 4]
        elif i == 6:
            return time_step * (x[:, 6]) + x[:, 5]
        elif i == 7:
            return time_step * (-10 * (x[:, 5] - x[:, 5] * x[:, 5] * x[:, 5] / 6.0) - x[:, 1]) + x[:, 6]
        elif i == 8:
            return time_step * (x[:, 8]) + x[:, 7]
        elif i == 9:
            return time_step * (-10 * (x[:, 7] - x[:, 7] * x[:, 7] * x[:, 7] / 6.0) - x[:, 1]) + x[:, 8]
        elif i == 10:
            return time_step * (x[:, 10]) + x[:, 9]
        elif i == 11:
            return time_step * (-10 * (x[:, 9] - x[:, 9] * x[:, 9] * x[:, 9] / 6.0) - x[:, 1]) + x[:, 10]
        elif i == 12:
            return time_step * (x[:, 12]) + x[:, 11]
        elif i == 13:
            return time_step * (-10 * (x[:, 11] - x[:, 11] * x[:, 11] * x[:, 11] / 6.0) - x[:, 1]) + x[:, 12]
        elif i == 14:
            return time_step * (x[:, 14]) + x[:, 13]
        elif i == 15:
            return time_step * (-10 * (x[:, 13] - x[:, 13] * x[:, 13] * x[:, 13] / 6.0) - x[:, 1]) + x[:, 14]
        elif i == 16:
            return time_step * (x[:, 16]) + x[:, 15]
        elif i == 17:
            return time_step * (-10 * (x[:, 15] - x[:, 15] * x[:, 15] * x[:, 15] / 6.0) - x[:, 1]) + x[:, 16]
        else:
            print("Vector function error!")
            exit()
    vf = torch.stack([f(i + 1, x) for i in range(DIM)], dim=1)
    return vf


# ############################################
# # for plotting
# ############################################
# PLOT_EXP_B = np.array([8, 8]) # sampling from domain for plotting the boundary of barrier using contour plot
# PLOT_LEN_B = np.power(2, PLOT_EXP_B) # the number of samples for each dimension of domain, usually larger than superp.DATA_LEN_D
#
# PLOT_EXP_V = np.array([7, 7]) # sampling from domain for plotting the vector field
# PLOT_LEN_V = np.power(2, PLOT_EXP_V) # the number of samples for each dimension of domain, usually equal to superp.DATA_LEN_D
#
# PLOT_EXP_P = np.array([7, 7]) # sampling from domain for plotting the scattering sampling points, should be equal to superp.DATA_LEN_D
# PLOT_LEN_P = np.power(2, PLOT_EXP_P) # the number of samples for each dimension of domain
#
# PLOT_VEC_SCALE = None